package com.llmops.core.node;

import com.llmops.core.Context;
import com.llmops.core.Node;
import com.llmops.core.NodeProperty;
import com.llmops.core.memory.MapDbMemoryStore;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.service.*;
import lombok.Data;

import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * 内容提取
 */
public class ExtractContentNode extends Node {
    @Data
    public static class LLMNodeProperty extends NodeProperty {

        private String apiUrl;
        private String apiKey;
        private String model;
        private Integer historyRecordCount;
        private String userPrompt;
        private String systemPrompt;
    }
    interface JsonExtractor {
        @UserMessage("""
            从文本中提取JSON格式信息,只返回JSON信息,不要输出多余的内容：
            文本：{{input}}
            要求:You must answer strictly in the following JSON format:{{jsonFormat}}
            """)
        String extract(@V("input") String text,@V("jsonFormat") String jsonFormat);
    }

    private LLMNodeProperty property = new LLMNodeProperty();

    public ExtractContentNode(String id,String title, Map<String, Object> config) {
        super(id,title, config);
        this.property.setApiUrl((String) config.getOrDefault("apiUrl", "https://api.deepseek.com/v1"));
        this.property.setApiKey((String) config.getOrDefault("apiKey","sk-0c3e6a9f9f7a4e63bb855290c544183c"));
        this.property.setModel((String) config.getOrDefault("model", "deepseek-chat"));
        this.property.setHistoryRecordCount((Integer) config.getOrDefault("historyRecordCount", 0));
        this.property.setUserPrompt((String) config.get("userPrompt"));
        this.property.setSystemPrompt((String) config.get("systemPrompt"));
    }

    @Override
    public void execute(Context context) {
        try {
            // 解析prompt中的变量
            String userPrompt;
            String systemPrompt;
            try {
                userPrompt = context.fillVariable(context,property.getUserPrompt());
                systemPrompt = context.fillVariable(context,property.getSystemPrompt());
            } catch (IllegalArgumentException e) {
                throw new RuntimeException("处理prompt变量时出错: " + e.getMessage(), e);
            }

            OpenAiChatModel chatModel = OpenAiChatModel.builder()
                    .baseUrl(property.getApiUrl())
                    .apiKey(property.getApiKey())
                    .modelName(property.getModel())
                    .logRequests(true)
                    .logResponses(true)
                    .build();

            JsonExtractor assistant = AiServices.builder(JsonExtractor.class)
                        .chatModel(chatModel)
                        .build();
            String answer = extractJsonBlock(assistant.extract(userPrompt,systemPrompt));
            context.setVariable(id + "_output", answer);
            logger.info(getId() + "_output={}", answer);
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException(String.format("节点={},运行异常",id));
        }
    }

    @Override
    public String getType() {
        return "extractContentNode";
    }

    private static final Pattern JSON_BLOCK_PATTERN = Pattern.compile("(?s)\\{.*\\}|\\[.*\\]");

    private String extractJsonBlock(String text) {
        Matcher matcher = JSON_BLOCK_PATTERN.matcher(text);
        if (matcher.find()) {
            return matcher.group();
        }
        return text;
    }
}